
"""
将YOLOv3的检测结果（Car.txt、Pedestrian.txt、Cyclist.txt）转换成json格式（COCO format）。
"""

import os
import json
from collections import OrderedDict

total_list = []

results_dir = 'E:\\0_My_codes\\ws_subl\\0217_yolov3_results\\'

labels = ['Car', 'Pedestrian', 'Cyclist']

label_to_cat = {
    "Car": 0,
    "Pedestrian": 1,
    "Cyclist": 2
}

for label in labels:
    # print(label)
    with open(results_dir + label + '.txt', "r") as f:
        for line in f.readlines():
            # print(line)
            line = line.strip().split(' ')
            # print(line)
            # transform---------------
            img_id = int(line[0])
            cat_id = label_to_cat[label]
            x = float("%.2f" % float(line[2]))
            y = float("%.2f" % float(line[3]))
            w = float("%.2f" % (float(line[4]) - float(line[2])))
            h = float("%.2f" % (float(line[5]) - float(line[3])))
            score = float(line[1])
            innt_dict = OrderedDict()
            innt_dict["image_id"] = img_id
            innt_dict["category_id"] = cat_id
            innt_dict["bbox"] = [x, y, w, h]
            innt_dict["score"] = score

            # innt_dict = {
            #     "image_id": img_id,
            #     "category_id": cat_id,
            #     "bbox": [x, y, w, h]
            # }
            total_list.append(innt_dict)
            # transform---------------

with open('test_2.json', 'w', encoding='utf-8') as f:
    json.dump(total_list, f, ensure_ascii=False)
